/*++

Module Name:

    simplewsk.c

Abstract:

    Wrapper library for WSK functions

Author:

    MaD, 12-May-2009

--*/

#include "WSKApi.h"

static WSK_REGISTRATION		g_WskRegistration;
static WSK_PROVIDER_NPI		g_WskProvider;
static WSK_CLIENT_DISPATCH	g_WskDispatch = {MAKE_WSK_VERSION(1,0), 0, NULL};

enum
{
	DEINITIALIZED,
	DEINITIALIZING,
	INITIALIZING,
	INITIALIZED
};

static LONG	g_SocketsState = DEINITIALIZED;

static
NTSTATUS
NTAPI
  CompletionRoutine(
    __in PDEVICE_OBJECT	DeviceObject,
    __in PIRP			Irp,
    __in PKEVENT		CompletionEvent
    )
{
	ASSERT( CompletionEvent );

    if( DeviceObject && Irp )
        ASSERT( CompletionEvent );

	KeSetEvent(CompletionEvent, IO_NO_INCREMENT, FALSE);
	return STATUS_MORE_PROCESSING_REQUIRED;
}

static
NTSTATUS
  InitWskData(
	__out PIRP*		pIrp,
	__out PKEVENT	CompletionEvent
	)
{
	ASSERT( pIrp );
	ASSERT( CompletionEvent );

	*pIrp = IoAllocateIrp(1, FALSE);
	if (!*pIrp) {
#if(DBG)
        DbgPrint("InitWskData(): IoAllocateIrp() failed\n");
#endif
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	KeInitializeEvent(CompletionEvent, SynchronizationEvent, FALSE);
	IoSetCompletionRoutine(*pIrp, CompletionRoutine, CompletionEvent, TRUE, TRUE, TRUE);
	return STATUS_SUCCESS;
}

static
NTSTATUS
  InitWskBuffer(
	__in  PVOID		Buffer,
	__in  ULONG		BufferSize,
	__out PWSK_BUF	WskBuffer
	)
{
	NTSTATUS Status = STATUS_SUCCESS;

	ASSERT( Buffer );
	ASSERT( BufferSize );
	ASSERT( WskBuffer );

	WskBuffer->Offset = 0;
	WskBuffer->Length = BufferSize;

	WskBuffer->Mdl = IoAllocateMdl(Buffer, BufferSize, FALSE, FALSE, NULL);
	if (!WskBuffer->Mdl) {
#if(DBG)
        DbgPrint("InitWskBuffer(): IoAllocateMdl() failed\n");
#endif
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	__try {
		MmProbeAndLockPages(WskBuffer->Mdl, KernelMode, IoWriteAccess);
	}
	__except(EXCEPTION_EXECUTE_HANDLER) {
#if(DBG)
        DbgPrint("InitWskBuffer(): MmProbeAndLockPages(%p) failed\n", Buffer);
#endif
		IoFreeMdl(WskBuffer->Mdl);
		Status = STATUS_ACCESS_VIOLATION;
	}

	return Status;
}

static
VOID
  FreeWskBuffer(
	__in PWSK_BUF WskBuffer
	)
{
	ASSERT( WskBuffer );

	MmUnlockPages(WskBuffer->Mdl);
	IoFreeMdl(WskBuffer->Mdl);
}

//
// Library initialization routine
//

NTSTATUS NTAPI SocketsInit()
{
	WSK_CLIENT_NPI	WskClient = {0};
	NTSTATUS		Status = STATUS_UNSUCCESSFUL;

	if (InterlockedCompareExchange(&g_SocketsState, INITIALIZING, DEINITIALIZED) != DEINITIALIZED)
		return STATUS_ALREADY_REGISTERED;

	WskClient.ClientContext = NULL;
	WskClient.Dispatch = &g_WskDispatch;

	Status = WskRegister(&WskClient, &g_WskRegistration);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("WskRegister() failed with status 0x%08X\n", Status);
#endif
		InterlockedExchange(&g_SocketsState, DEINITIALIZED);
		return Status;
	}

	Status = WskCaptureProviderNPI(&g_WskRegistration, WSK_INFINITE_WAIT, &g_WskProvider);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("WskCaptureProviderNPI() failed with status 0x%08X\n", Status);
#endif
		WskDeregister(&g_WskRegistration);
		InterlockedExchange(&g_SocketsState, DEINITIALIZED);
		return Status;
	}

	InterlockedExchange(&g_SocketsState, INITIALIZED);
	return STATUS_SUCCESS;
}

//
// Library deinitialization routine
//

VOID NTAPI SocketsDeinit()
{
	if (InterlockedCompareExchange(&g_SocketsState, INITIALIZED, DEINITIALIZING) != INITIALIZED)
		return;

	WskReleaseProviderNPI(&g_WskRegistration);
	WskDeregister(&g_WskRegistration);

	InterlockedExchange(&g_SocketsState, DEINITIALIZED);
}



PWSK_SOCKET
NTAPI
  CreateSocket(
    __in ADDRESS_FAMILY	AddressFamily,
    __in USHORT			SocketType,
    __in ULONG			Protocol,
    __in ULONG			Flags
    )
{
	KEVENT			CompletionEvent = {0};
	PIRP			Irp = NULL;
	PWSK_SOCKET		WskSocket = NULL;
	NTSTATUS		Status = STATUS_UNSUCCESSFUL;

	if (g_SocketsState != INITIALIZED)
		return NULL;

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("CreateSocket(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return NULL;
	}

	Status = g_WskProvider.Dispatch->WskSocket(
		g_WskProvider.Client,
		AddressFamily,
		SocketType,
		Protocol,
		Flags,
		NULL,
		NULL,
		NULL,
		NULL,
		NULL,
		Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}
	
	WskSocket = NT_SUCCESS(Status) ? (PWSK_SOCKET)Irp->IoStatus.Information : NULL;

	IoFreeIrp(Irp);
	return (PWSK_SOCKET)WskSocket;
}

NTSTATUS
NTAPI
  CloseSocket(
	__in PWSK_SOCKET WskSocket
	)
{
	KEVENT		CompletionEvent = {0};
	PIRP		Irp = NULL;
	NTSTATUS	Status = STATUS_UNSUCCESSFUL;

	if (g_SocketsState != INITIALIZED || !WskSocket)
		return STATUS_INVALID_PARAMETER;

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("CloseSocket(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return Status;
	}

	Status = ((PWSK_PROVIDER_BASIC_DISPATCH)WskSocket->Dispatch)->WskCloseSocket(WskSocket, Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}

	IoFreeIrp(Irp);
	return Status;
}


NTSTATUS
NTAPI
  Connect(
	__in PWSK_SOCKET	WskSocket,
	__in PSOCKADDR		RemoteAddress
	)
{
	KEVENT		CompletionEvent = {0};
	PIRP		Irp = NULL;
	NTSTATUS	Status = STATUS_UNSUCCESSFUL;

	if (g_SocketsState != INITIALIZED || !WskSocket || !RemoteAddress)
		return STATUS_INVALID_PARAMETER;

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("Connect(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return Status;
	}

	Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)->WskConnect(
		WskSocket,
		RemoteAddress,
		0,
		Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}

	IoFreeIrp(Irp);
	return Status;
}

PWSK_SOCKET
NTAPI
  SocketConnect(
	__in USHORT		SocketType,
	__in ULONG		Protocol,
	__in PSOCKADDR	RemoteAddress,
	__in PSOCKADDR	LocalAddress
	)
{
	KEVENT			CompletionEvent = {0};
	PIRP			Irp = NULL;
	NTSTATUS		Status = STATUS_UNSUCCESSFUL;
	PWSK_SOCKET		WskSocket = NULL;

	if (g_SocketsState != INITIALIZED || !RemoteAddress || !LocalAddress)
		return NULL;

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("InitWskData() failed with status 0x%08X\n", Status);
#endif
		return NULL;
	}

	Status = g_WskProvider.Dispatch->WskSocketConnect(
		g_WskProvider.Client,
		SocketType,
		Protocol,
		LocalAddress,
		RemoteAddress,
		0,
		NULL,
		NULL,
		NULL,
		NULL,
		NULL,
		Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}

	WskSocket = NT_SUCCESS(Status) ? (PWSK_SOCKET)Irp->IoStatus.Information : NULL;

	IoFreeIrp(Irp);
	return WskSocket;
}


LONG
NTAPI
  Send(
	__in PWSK_SOCKET	WskSocket,
	__in PVOID			Buffer,
	__in ULONG			BufferSize,
	__in ULONG			Flags
	)
{
	KEVENT		CompletionEvent = {0};
	PIRP		Irp = NULL;
	WSK_BUF		WskBuffer = {0};
	LONG		BytesSent = SOCKET_ERROR;
	NTSTATUS	Status = STATUS_UNSUCCESSFUL;

	if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize)
		return SOCKET_ERROR;

	Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("Send(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return SOCKET_ERROR;
	}

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("Send(): InitWskData() failed with status 0x%08X\n", Status);
#endif
        FreeWskBuffer(&WskBuffer);
		return SOCKET_ERROR;
	}

	Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)->WskSend(
		WskSocket,
		&WskBuffer,
		Flags,
		Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}
	
	BytesSent = NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;

	IoFreeIrp(Irp);
	FreeWskBuffer(&WskBuffer);
	return BytesSent;
}

LONG
NTAPI
  SendTo(
	__in PWSK_SOCKET	WskSocket,
	__in PVOID			Buffer,
	__in ULONG			BufferSize,
    __in_opt PSOCKADDR	RemoteAddress
	)
{
	KEVENT		CompletionEvent = {0};
	PIRP		Irp = NULL;
	WSK_BUF		WskBuffer = {0};
	LONG		BytesSent = SOCKET_ERROR;
	NTSTATUS	Status = STATUS_UNSUCCESSFUL;

	if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize)
		return SOCKET_ERROR;

	Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("SendTo(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return SOCKET_ERROR;
	}

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("SendTo(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		FreeWskBuffer(&WskBuffer);
		return SOCKET_ERROR;
	}

	Status = ((PWSK_PROVIDER_DATAGRAM_DISPATCH)WskSocket->Dispatch)->WskSendTo(
		WskSocket,
		&WskBuffer,
		0,
		RemoteAddress,
		0,
		NULL,
		Irp);
	if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);//&LTimeOut);
		Status = Irp->IoStatus.Status;
	}

    if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("SendTo(): FAILED NT_STATUS = 0x%X\n", Status);
#endif
    }

    BytesSent = NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;

	IoFreeIrp(Irp);
	FreeWskBuffer(&WskBuffer);
	return BytesSent;
}

LONG
NTAPI
  Receive(
	__in  PWSK_SOCKET	WskSocket,
	__out PVOID			Buffer,
	__in  ULONG			BufferSize,
	__in  ULONG			Flags
	)
{
	KEVENT		CompletionEvent = {0};
	PIRP		Irp = NULL;
	WSK_BUF		WskBuffer = {0};
	LONG		BytesReceived = SOCKET_ERROR;
	NTSTATUS	Status = STATUS_UNSUCCESSFUL;

	if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize)
		return SOCKET_ERROR;

	Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("Receive(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return SOCKET_ERROR;
	}

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("Receive(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		FreeWskBuffer(&WskBuffer);
		return SOCKET_ERROR;
	}

	Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)->WskReceive(
		WskSocket,
		&WskBuffer,
		Flags,
		Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}
	
	BytesReceived = NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;

	IoFreeIrp(Irp);
	FreeWskBuffer(&WskBuffer);
	return BytesReceived;
}

LONG
NTAPI
  ReceiveFrom(
	__in  PWSK_SOCKET	WskSocket,
	__out PVOID			Buffer,
	__in  ULONG			BufferSize,
	__out_opt PSOCKADDR	RemoteAddress,
    __out_opt PULONG	ControlFlags
	)
{
	KEVENT		CompletionEvent = {0};
	PIRP		Irp = NULL;
	WSK_BUF		WskBuffer = {0};
	LONG		BytesReceived = SOCKET_ERROR;
	NTSTATUS	Status = STATUS_UNSUCCESSFUL;

    if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize){
    #if(DBG)
            MpPrintf("ReceiveFrom(): FAILED NT_STATUS =  0x%X",Status);
    #endif
		return SOCKET_ERROR;
    }
	Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("ReceiveFrom(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return SOCKET_ERROR;
	}

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("ReceiveFrom(): InitWskData() failed NT_STATUS =  0x%08X\n", Status);
#endif
		FreeWskBuffer(&WskBuffer);
		return SOCKET_ERROR;
	}

	Status = ((PWSK_PROVIDER_DATAGRAM_DISPATCH)WskSocket->Dispatch)->WskReceiveFrom(
		WskSocket,
		&WskBuffer,
		0,
		RemoteAddress,
		0,
		NULL,
		ControlFlags,
		Irp);
	if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL );//&LTimeOut);
        Status = Irp->IoStatus.Status;
	}

    if (!NT_SUCCESS(Status)){
#if(DBG)
        MpPrintf("ReceiveFrom(): FAILED NT_STATUS =  0x%X",Status);
#endif

    }

    BytesReceived = NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;

    IoFreeIrp(Irp);
	FreeWskBuffer(&WskBuffer);
	return BytesReceived;

}

NTSTATUS
NTAPI
  Bind(
	__in PWSK_SOCKET	WskSocket,
	__in PSOCKADDR		LocalAddress
	)
{
	KEVENT		CompletionEvent = {0};
	PIRP		Irp = NULL;
	NTSTATUS	Status = STATUS_UNSUCCESSFUL;

	if (g_SocketsState != INITIALIZED || !WskSocket || !LocalAddress)
		return STATUS_INVALID_PARAMETER;

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("Bind(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return Status;
	}

	Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)->WskBind(
		WskSocket,
		LocalAddress,
		0,
		Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}

	IoFreeIrp(Irp);
	return Status;
}


PWSK_SOCKET
NTAPI
  Accept(
	__in PWSK_SOCKET	WskSocket,
	__out_opt PSOCKADDR	LocalAddress,
	__out_opt PSOCKADDR	RemoteAddress
   )
{
	KEVENT			CompletionEvent = {0};
	PIRP			Irp = NULL;
	NTSTATUS		Status = STATUS_UNSUCCESSFUL;
	PWSK_SOCKET		AcceptedSocket = NULL;

	if (g_SocketsState != INITIALIZED || !WskSocket)
		return NULL;

	Status = InitWskData(&Irp, &CompletionEvent);
	if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("Accept(): InitWskData() failed with status 0x%08X\n", Status);
#endif
		return NULL;
	}

	Status = ((PWSK_PROVIDER_LISTEN_DISPATCH)WskSocket->Dispatch)->WskAccept(
		WskSocket,
		0,
		NULL,
		NULL,
		LocalAddress,
		RemoteAddress,
		Irp);
	if (Status == STATUS_PENDING) {
		KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
		Status = Irp->IoStatus.Status;
	}

	AcceptedSocket = NT_SUCCESS(Status) ? (PWSK_SOCKET)Irp->IoStatus.Information : NULL;

	IoFreeIrp(Irp);
	return AcceptedSocket;
}


NTSTATUS MyWskReceiveFromEvent(
        PVOID SocketContext,
        unsigned long Flags,
        PWSK_DATAGRAM_INDICATION DataIndication);

FAST_MUTEX gSocketMutex;

static PWSK_SOCKET		gLocalSocket = NULL;

const WSK_CLIENT_DATAGRAM_DISPATCH WskClientReceiveDispatch = {
   (PFN_WSK_RECEIVE_FROM_EVENT)MyWskReceiveFromEvent
};


NTSTATUS InitSocketsExt() {

    NTSTATUS status;
    WSK_CLIENT_NPI wskClientNpi = {0};

    //PIRP SocketIrp;
    SOCKADDR_IN		gServerAddr = {0};

    DbgPrint(("===>InitSockets...\n"));

    wskClientNpi.ClientContext = NULL;
    wskClientNpi.Dispatch = &g_WskDispatch;
    status = WskRegister(&wskClientNpi, &g_WskRegistration);

    __debugbreak();

    if(!NT_SUCCESS(status)) {
        DbgPrint("InitSockets: ERROR WskRegister - %08x\n", status);
        return status;
    } // Capture the WSK Provider NPI. If WSK subsystem is not ready yet, // wait until it becomes ready.

    status = WskCaptureProviderNPI(&g_WskRegistration, WSK_INFINITE_WAIT, &g_WskProvider);

    if(!NT_SUCCESS(status)) {
        DbgPrint("InitWorker: ERROR WskCaptureProviderNPI - %08x\n", status);
        WskDeregister(&g_WskRegistration);
        return status;
    }

    gServerAddr.sin_family = AF_INET;
    gServerAddr.sin_port = HTONS(40007);
    gServerAddr.sin_addr.s_addr	= INADDR_ANY;

    ExInitializeFastMutex(&gSocketMutex); //set up the socket

    CreateSocketExt();
    DbgPrint(("<===InitSockets...\n"));

    return STATUS_SUCCESS;
}

NTSTATUS MyWskReceiveFromEvent(
        PVOID SocketContext,
        unsigned long Flags,
        PWSK_DATAGRAM_INDICATION DataIndication)
{
    //NTSTATUS ntStatus;
    //unsigned int totalPacketSize;
    //wchar_t * wcharArray;
    //KdPrint(("===>MyWskReceiveFromEvent...\n"));
    //SendPacketWSK(gLocalSocket, &gServerAddr, "WskRec", 6, NULL);

    PWSK_BUF pwsk_buff;

    PETHERNET pEthePack;

    PMDL pMdl;

    PVOID pBuff;

    ULONG MdlLen;

    PSOCKADDR pRemoteAddress;

    ULONG BuffLen, Offset;

    if(!DataIndication){
        DbgPrint("MyWskReceiveFromEvent: null DataIndication");
        return STATUS_DATA_NOT_ACCEPTED;
    }
    if(DataIndication->Buffer.Length == 0 || DataIndication->Buffer.Mdl == NULL){
        DbgPrint("MyWskReceiveFromEvent: invalid DataIndication");
    }

    while(DataIndication != NULL){

        pwsk_buff = &DataIndication->Buffer;
        pRemoteAddress = DataIndication->RemoteAddress;
        pMdl = pwsk_buff->Mdl;

        //pBuff = (PVOID)MmGetSystemAddressForMdlSafe((PMDL)pMdl,HighPagePriority );

        pEthePack = (PETHERNET)MmGetMdlVirtualAddress(pMdl );

        MdlLen = MmGetMdlByteCount(pMdl);
        Offset = MmGetMdlByteOffset(pMdl);

        if ((pEthePack != NULL)&&( MdlLen > sizeof(ETHERNET) )) {

            pBuff = ((unsigned char*)pEthePack) + sizeof(ETHERNET);

            BuffLen = MdlLen - pEthePack->udp_head.Length;


        }




//        if( ( pEthePack->IPType == 0x0800 ) &&
//                (((pEthePack->ip_head.Protocol & 0xFF0000)>>16) == 0x11) &&
//                    ()){



//        }

        DataIndication = DataIndication->Next;
        KdPrint(("ReceiveFromEvent incoming package \n"));
    }

    //Flags = 0; // TEST

    if((SocketContext)||(Flags))
        DbgPrint("SocketContext ==0; Flags == 0");

    //SendPacketWSK(gLocalSocket, &gServerAddr, "WskRec", 6, NULL);
    //KdPrint(("<===MyWskReceiveFromEvent...\n"));
    return STATUS_SUCCESS;
}


NTSTATUS CreateSocketExt() {
    NTSTATUS status;
    PIRP SocketIrp;
    WSK_EVENT_CALLBACK_CONTROL callbackControl;
    DbgPrint(("===>CreateSocket...\n"));
    SocketIrp = IoAllocateIrp(1,FALSE);

    if (!SocketIrp) {

        DbgPrint("CreateSocket: ERROR IoAllocateIrp failed\n");
        return STATUS_INSUFFICIENT_RESOURCES;

    } // First, configure the WSK client such that WskReceiveFromEvent callback // will be automatically enabled on the listening socket upon creation.

    callbackControl.NpiId = (PNPIID)&NPI_WSK_INTERFACE_ID;
    callbackControl.EventMask = WSK_EVENT_RECEIVE_FROM;

    status = g_WskProvider.Dispatch->WskControlClient(
                                                        g_WskProvider.Client,
                                                        WSK_SET_STATIC_EVENT_CALLBACKS,
                                                        sizeof(WSK_EVENT_CALLBACK_CONTROL),
                                                        &callbackControl,
                                                        0,
                                                        NULL,
                                                        NULL,
                                                        NULL
                                                      );

    if(!NT_SUCCESS(status)) {

        DbgPrint("CreateSocket: WskControlClient failed\n");
        return STATUS_UNSUCCESSFUL;

    }

    IoSetCompletionRoutine(	SocketIrp,
                            SocketCompletionExt,
                            NULL,
                            TRUE,
                            TRUE,
                            TRUE);

    // Initiate the creation of the socket
    status = g_WskProvider.Dispatch->WskSocket(g_WskProvider.Client,
                                                AF_INET,
                                                SOCK_DGRAM,
                                                IPPROTO_UDP,
                                                WSK_FLAG_DATAGRAM_SOCKET,
                                                NULL,
                                                &WskClientReceiveDispatch,
                                                NULL,
                                                NULL,
                                                NULL,
                                                SocketIrp);
    if(!NT_SUCCESS(status)) {
        DbgPrint("CreateSocket: ERROR WskSocket - %08x\n", status);
        return STATUS_UNSUCCESSFUL;
    }
    DbgPrint(("<===CreateSocket...\n")); return status;
}

NTSTATUS SocketCompletionExt (
                                PDEVICE_OBJECT DeviceObject,
                                PIRP Irp,
                                PVOID Context
                             ) {

    SOCKADDR_IN		local_addr = {0};

    NTSTATUS ntStatus;
    PWSK_SOCKET socket;
    PIRP BindIrp;
    UNREFERENCED_PARAMETER(DeviceObject);
    DbgPrint(("===>SocketCompletion...\n")); // Check the result of the socket creation

    if((DeviceObject==0) || (Context == 0) )
        DbgPrint(("===>DeviceObject == 0...\n"));

    if (Irp->IoStatus.Status != STATUS_SUCCESS) {

        IoFreeIrp(Irp);
        return STATUS_UNSUCCESSFUL;

    }
    // Save the socket object for the new socket

    socket = (PWSK_SOCKET)(Irp->IoStatus.Information);
    // First we need to bind the socket to a local address
    local_addr.sin_family = AF_INET;
    local_addr.sin_port = HTONS(40017);
    local_addr.sin_addr.S_un.S_addr = INADDR_ANY;
    memset(local_addr.sin_zero, 0, sizeof(local_addr.sin_zero));

    BindIrp = IoAllocateIrp(1,FALSE);

    if (!BindIrp) {
        DbgPrint("SocketCompletion: ERROR IoAllocateIrp failed\n");
        IoReuseIrp(Irp, STATUS_UNSUCCESSFUL);
        //IoSetCompletionRoutine(Irp, FreeIrpCompletion, NULL, TRUE, TRUE, TRUE);
        ((PWSK_PROVIDER_BASIC_DISPATCH)socket->Dispatch)->WskCloseSocket(socket, Irp);
        ExAcquireFastMutex(&gSocketMutex);
        gLocalSocket = NULL;
        ExReleaseFastMutex(&gSocketMutex);
        __debugbreak();
        return STATUS_UNSUCCESSFUL;
    } //notice we pass socket on to the BindCompletion routine

    IoSetCompletionRoutine(BindIrp, BindCompletionExt, socket, TRUE, TRUE, TRUE);

    ntStatus = ((PWSK_PROVIDER_DATAGRAM_DISPATCH)socket->Dispatch)->WskBind(socket, (PSOCKADDR)&local_addr, 0, BindIrp);

    if (!NT_SUCCESS(ntStatus)) {
        DbgPrint("SocketCompletion: WskBindSocket failed - %08x\n", ntStatus);
        IoReuseIrp(Irp, STATUS_UNSUCCESSFUL);
        //IoSetCompletionRoutine(Irp, FreeIrpCompletion, NULL, TRUE, TRUE, TRUE);
        ((PWSK_PROVIDER_BASIC_DISPATCH)socket->Dispatch)->WskCloseSocket(socket, Irp);
        ExAcquireFastMutex(&gSocketMutex);
        gLocalSocket = NULL;
        ExReleaseFastMutex(&gSocketMutex);
        __debugbreak();
        return STATUS_UNSUCCESSFUL;
    }

    IoFreeIrp(Irp);
    DbgPrint(("<===SocketCompletion...\n"));
    return STATUS_MORE_PROCESSING_REQUIRED;
}



NTSTATUS BindCompletionExt(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context) {
    //NTSTATUS ntStatus;
    //PIRP SetRemoteIrp;
    PWSK_SOCKET socket = (PWSK_SOCKET)Context;
    PWSK_PROVIDER_CONNECTION_DISPATCH Dispatch;

    if((!DeviceObject))
        DbgPrint(("DeviceObject == 0; SetRemoteIrp == 0"));

    DbgPrint(("==>BindCompletion\n")); // Check the result of binding the local address
    if (Irp->IoStatus.Status != STATUS_SUCCESS) {
        DbgPrint("BindCompletion: Failed to bind local address - %08x\n", Irp->IoStatus.Status);
        IoReuseIrp(Irp, STATUS_UNSUCCESSFUL);
        //IoSetCompletionRoutine(Irp, FreeIrpCompletion, NULL, TRUE, TRUE, TRUE);
        ((PWSK_PROVIDER_BASIC_DISPATCH)socket->Dispatch)->WskCloseSocket(socket, Irp);
        ExAcquireFastMutex(&gSocketMutex);
        gLocalSocket = NULL;
        ExReleaseFastMutex(&gSocketMutex);
        __debugbreak();
        return STATUS_MORE_PROCESSING_REQUIRED;
    }

    ExAcquireFastMutex(&gSocketMutex);
    gLocalSocket = socket;
    ExReleaseFastMutex(&gSocketMutex); // Get pointer to the socket's provider dispatch structure

    Dispatch = (PWSK_PROVIDER_CONNECTION_DISPATCH)(socket->Dispatch);
    IoFreeIrp(Irp);
    DbgPrint(("<==BindCompletion\n"));
    return STATUS_MORE_PROCESSING_REQUIRED;
}

NTSTATUS
NTAPI
  CloseSocketExt(
    __in PWSK_SOCKET WskSocket
    )
{
    KEVENT		CompletionEvent = {0};
    PIRP		Irp = NULL;
    NTSTATUS	Status = STATUS_UNSUCCESSFUL;

    Status = InitWskData(&Irp, &CompletionEvent);
    if (!NT_SUCCESS(Status)) {
#if(DBG)
        DbgPrint("CloseSocket(): InitWskData() failed with status 0x%08X\n", Status);
#endif
        return Status;
    }

    Status = ((PWSK_PROVIDER_BASIC_DISPATCH)WskSocket->Dispatch)->WskCloseSocket(WskSocket, Irp);
    if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
        Status = Irp->IoStatus.Status;
    }

    IoFreeIrp(Irp);
    return Status;
}


VOID NTAPI DeInitSocketExt()
{

    if(gLocalSocket){

        CloseSocketExt(gLocalSocket);
        gLocalSocket = NULL;

    }

    WskReleaseProviderNPI(&g_WskRegistration);
    WskDeregister(&g_WskRegistration);

}


/*
NTSTATUS MyWskReceiveFromEvent(PVOID SocketContext, unsigned long Flags, PWSK_DATAGRAM_INDICATION DataIndication) {
    NTSTATUS ntStatus;
    unsigned int totalPacketSize;
    wchar_t * wcharArray;
    DbgPrint(("===>MyWskReceiveFromEvent...\n"));
    SendPacketWSK(gLocalSocket, &gServerAddr, "WskRec", 6, NULL);
    if(!DataIndication){
        DbgPrint("MyWskReceiveFromEvent: null DataIndication");
        return STATUS_DATA_NOT_ACCEPTED;
    }
    if(DataIndication->Buffer.Length == 0 || DataIndication->Buffer.Mdl == NULL){
        DbgPrint("MyWskReceiveFromEvent: invalid DataIndication");
    }
    SendPacketWSK(gLocalSocket, &gServerAddr, "WskRec", 6, NULL);
    DbgPrint(("<===MyWskReceiveFromEvent...\n"));
    return STATUS_SUCCESS;
}*/



